import heapq
import time
from collections import deque

import numpy as np


def bfs(grid, a, b, c, d):
    m, n = len(grid), len(grid[0])
    q = deque([(0, a, b)])
    visited = np.full((m, n), False)
    visited[a][b] = True
    while q:
        for _ in range(len(q)):
            cost, x, y = q.popleft()
            if x == c and y == d:
                return cost
            for tx, ty in [(x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)]:
                if 0 <= tx < m and 0 <= ty < n and (not visited[tx][ty]) and grid[tx][ty] == 1:
                    visited[tx][ty] = True
                    q.append((cost + 1, tx, ty))
    return -1


def dijkstra(grid, a, b, c, d):
    m, n = len(grid), len(grid[0])
    visited = np.full((m, n), False)
    distances = np.full((m, n), np.inf)
    distances[a, b] = 0
    heap = [(0, a, b)]
    while heap:
        cost, x, y = heapq.heappop(heap)
        if visited[x][y]: continue
        visited[x][y] = True
        if x == c and y == d:
            return distances[x][y]
        for tx, ty in [(x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)]:
            if 0 <= tx < m and 0 <= ty < n and grid[tx][ty] == 1 and distances[x][y] + 1 < distances[tx][ty]:
                distances[tx][ty] = distances[x][y] + 1
                heapq.heappush(heap, (distances[tx][ty], tx, ty))
    return -1

# 曼哈顿距离L1 norm
def manhattan_distance(a, b, c, d):
    return abs(a - c) + abs(b - d)
# 切比雪夫距离 L2norm(对角线距离)
def chebyshev_distance(a, b, c, d):
    return max(abs(a - c), abs(b - d))
# 欧几里得距离
def euclidean_distance(a, b, c, d):
    # return ((a-c) ** 2 + (b-d) ** 2) ** 0.5
    return (a-c) ** 2 + (b-d) ** 2


def astar(grid, a, b, c, d):
    m, n = len(grid), len(grid[0])
    visited = np.full((m, n), False)
    distances = np.full((m, n), np.inf)
    distances[a][b] = 0
    heap = [(0, a, b)]
    while heap:
        cost, x, y = heapq.heappop(heap)
        if visited[x][y]: continue
        visited[x][y] = True
        if x == c and y == d:
            return distances[x][y]
        for tx, ty in [(x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)]:
            if 0 <= tx < m and 0 <= ty < n and grid[tx][ty] == 1 and distances[x][y] + 1 < distances[tx][ty]:
                distances[tx][ty] = distances[x][y] + 1
                # heapq.heappush(heap, (distances[tx][ty] + manhattan_distance(tx, ty, c, d), tx, ty))
                # heapq.heappush(heap, (distances[tx][ty] + chebyshev_distance(tx, ty, c, d), tx, ty))
                heapq.heappush(heap, (distances[tx][ty] + euclidean_distance(tx, ty, c, d), tx, ty))
    return -1


def test(epoches, maximum_n):
    for epoch in range(1, epoches + 1):
        m, n = np.random.randint(1, maximum_n + 1), np.random.randint(1, maximum_n + 1)
        grid = np.random.randint(0, 2, size=(m, n))
        x_start, y_start = 0, 0
        x_target, y_target = np.random.randint(0, m), np.random.randint(0, n)
        grid[x_start, y_start] = grid[x_target, y_target] = 1
        s1 = time.time()
        r1 = bfs(grid, x_start, y_start, x_target, y_target)
        s2 = time.time()
        r2 = dijkstra(grid, x_start, y_start, x_target, y_target)
        s3 = time.time()
        r3 = astar(grid, x_start, y_start, x_target, y_target)
        s4 = time.time()
        ans = r1 == r2 == r3
        print(f'epoche: {epoch}: 正确: {ans}, bfs耗时: {(s2 - s1) * 1000:5f}ms, dijkstra耗时：{(s3 - s2) * 1000:5f}ms, astar耗时: {(s4 - s3) * 1000:5f}ms')
        # print(f'epoche: {epoch}: grid.shape: {grid.shape}, 正确: {ans}, bfs耗时: {s2 - s1:5f}s, dijkstra耗时：{s3 - s2:5f}s, astar耗时: {s4 - s3:5f}s')


if __name__ == "__main__":
    maximum_n = 5000
    epoches = 100
    test(epoches, maximum_n)
